import os
from typing import Any

import pandas as pd
import numpy as np
import rasterio
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm

from tif_stats import wp_cost_type_mapping, tco2_cost_type_mapping

# 行政编码到中文名的映射
admin_code_to_name_default = {
    "110000": "北京市",
    "120000": "天津市",
    "130000": "河北省",
    "140000": "山西省",
    "150000": "内蒙古自治区",
    "210000": "辽宁省",
    "220000": "吉林省",
    "230000": "黑龙江省",
    "310000": "上海市",
    "320000": "江苏省",
    "330000": "浙江省",
    "340000": "安徽省",
    "350000": "福建省",
    "360000": "江西省",
    "370000": "山东省",
    "410000": "河南省",
    "420000": "湖北省",
    "430000": "湖南省",
    "440000": "广东省",
    "450000": "广西壮族自治区",
    "460000": "海南省",
    "500000": "重庆市",
    "510000": "四川省",
    "520000": "贵州省",
    "530000": "云南省",
    "540000": "西藏自治区",
    "610000": "陕西省",
    "620000": "甘肃省",
    "630000": "青海省",
    "640000": "宁夏回族自治区",
    "650000": "新疆维吾尔自治区",
    "660000": "新疆建设兵团",
}

def get_excel_value_by_key(excel_path: str, key_column: str, key: int, value_column: str) -> Any | None:
    """
    通过指定的键（列名1的值）来获取另一列（列名2）的值。

    :param excel_path: Excel文件的路径
    :param key: 指定列的值
    :param key_column: 指定列的列名
    :param value_column: 要获取值的列名
    :return: 二氧化硫治理成本，如果未找到则返回None
    """
    try:
        # 读取Excel文件
        df = pd.read_excel(excel_path)

        # 查找对应行政区划代码的行
        result = df[df[key_column] == key]
        if not result.empty:
            # 返回治理成本
            return result[value_column].values[0]
        else:
            # 如果未找到对应行政区划代码，返回None
            return None
    except Exception as e:
        print(f"发生错误: {e}")
        return None

def process_tif_file(file_path: str, folder_path: str, excel_path: str, cost_type_mapping: dict, admin_code_to_name: dict) -> dict:
    """
    处理单个 TIFF 文件的函数
    返回一个字典，包含文件名和一些统计信息
    """
    try:
        filename = os.path.basename(file_path)
        admin_code = int(filename.split('_')[0])
        quantity_type = filename.split('_')[1]
        ecosystem_type = filename.split('_')[2]

        with rasterio.open(file_path) as src:
            data = src.read(1)

            Q_total_value = np.nansum(data)  # 忽略NaN值
            Q_mean_value = np.nanmean(data) / 0.0009  # 实物量地均值（t/km2）

            area = np.sum(np.logical_not(np.isnan(data))) * 0.0009  # 面积（km2）
            self_mean_value = Q_total_value / area

            V_cost = get_excel_value_by_key(excel_path, '行政区划代码', admin_code, cost_type_mapping.get(quantity_type))
            V_total_value = Q_total_value * V_cost
            V_mean_value = V_total_value / area

        return {
            '行政区划代码': admin_code,
            '核算地区': admin_code_to_name.get(str(admin_code), '未知地区'),
            '面积（km2）': area,
            '调节服务项': quantity_type,
            '生态系统': ecosystem_type,
            '实物量总值（t）': Q_total_value,
            '实物量地均值（t/km2）': Q_mean_value,
            '价格（元/t）': V_cost,
            '价值量总值（元）': V_total_value,
            '价值量地均值（元/km2）': V_mean_value
        }
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {e}")
        return None

def analyze_tif_files(folder_path: str, excel_path: str, cost_type_mapping: dict, admin_code_to_name: dict = None, output_path: str = None) -> None:
    """
    分析指定文件夹中的 TIF 文件，并将结果保存到 Excel 文件中。根据 tif 文件名的信息，制作列名。

    :param folder_path: 包含 TIF 文件的文件夹路径
    :param excel_path: Excel 文件路径，用于查找价格
    :param cost_type_mapping: 调节服务项到价格列名的映射字典
    :param admin_code_to_name: 行政编码到中文名的映射字典
    :param output_path: 结果保存的 Excel 文件路径，默认为 None，将根据 folder_path 生成
    """
    if admin_code_to_name is None:
        admin_code_to_name = admin_code_to_name_default

    # 定义列名
    columns = ['行政区划代码', '核算地区', '面积（km2）', '调节服务项', '生态系统', '实物量总值（t）', '实物量地均值（t/km2）', '价格（元/t）', '价值量总值（元）', '价值量地均值（元/km2）']

    # 获取文件夹中的所有 TIF 文件
    tif_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.tif')]

    # 使用多进程处理每个文件
    results = []
    with ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
        futures = [executor.submit(process_tif_file, file_path, folder_path, excel_path, cost_type_mapping, admin_code_to_name) for file_path in tif_files]
        for future in tqdm(as_completed(futures), total=len(futures), desc="Processing TIF files"):
            result = future.result()
            if result is not None:
                results.append(result)

    # 将结果转换为 DataFrame
    df = pd.DataFrame(results, columns=columns)

    # 根据 folder_path 生成默认的 output_path
    if output_path is None:
        folder_name = os.path.basename(folder_path)
        parent_folder_path = os.path.dirname(folder_path)
        output_path = os.path.join(parent_folder_path, f'{folder_name}_results.xlsx')

    # 将结果存储到 Excel 文件
    df.to_excel(output_path, index=False)

    print(f'结果已保存到 {output_path}')

# 示例主函数
def main():
    folder_path = r'H:\Qtco2'
    excel_path = r'H:\GEP_excel_data\V_price\价格表-总表.xlsx'

    analyze_tif_files(folder_path, excel_path, tco2_cost_type_mapping)

if __name__ == '__main__':
    main()